I think this should be named better. Perhaps split_attn_cudagraph? I also don't understand why this has to be a flag and we can't just ask the attention backend what it wants?

I think we must leave such a flag in the global config, which tells the compiler backend to do the right thing. Otherwise, how is the attention backend supposed to communicate its requirements to the compiler? At least for now, the force_separate_routine flag of an attention backend has the ability to enforce its preference during the initialize_attn_backend phase of the gpu model runner.

I think this should be named better. Perhaps split_attn_cudagraph?

I am not sure what name can be better. Btw, I'm afraid split_attn_cudagraph is not a good name. It sounds like splitting the full graph into be piecewise graph, where attn ops are the splitting ops, like what we have already done.

Good call on the name. Also makes sense we use this to communicate from attention backend to compiler. Let's make sure that happens inside set_splitting_ops_for_v1/somewhere inside config initialization, if we can.

I think we should figure out a different name for this; the current name doesnt indicate any relation to cudagraphs

I'm not as zoned into this PR as you folks are, but I have no clue what this flag is from the name.

I think we should figure out a different name for this; the current name doesnt indicate any relation to cudagraphs

How about cudagraph_separate_routine? Cutting the "attention" out seems to have no effect on its meaning. While it is basically prepared for distinct attention routines that are actually executed, in the future, that may be more than just attention ops.

Just changed to cudagraph_separate_routine. It should be better.

Really nice explanation and the logic is clear, thanks!

nit: maybe we can rename this file in a future PR (that just does file renaming) so we can see the diff here?

Done!

nit: I think we should name this use_full_cudagraph I could see cases in the future where we might want to mixed piecewise + full-cudagraphs but have a different heuristic than is_pure_decode. e.g. for FA3 we may want to run mixed small decodes or spec-decode using full-cudagraphs but run large prefills using piecewise cudagraphs

is pure decode is not equivalent to full cudagraphs. It is the same for backends that only support pure decode batches in CG, but not otherwise. While I agree the name could be improved I don't think full_cudagraph is better

ya no I agree; wrote this early in the review, sorry!. Im actually more of fan of uniform_batch now haha, see: #20059 (review)

Just updated to uniform_batch (along with some related logic), though I haven't yet tested if it's okay with speculative decoding.

ditto above; I dont think we over index on is_pure_decode naming and since we may want to dispatch between full-cudagraphs and piecewise on other metrics in the future

Based on this description I think the warning is very confusing to the user. They won't know what a "separate attention routine" means. We should try to think of new naming here; I would prefer something like:

{attn_backend_i.__name__} generally performs better when using
full cuda-graphs for only pure decode batches (and using piecewise
cuda-graphs for prefill and mixed-prefill-decode iterations) to enable
this behavior turn on 
CompilationConfig.mixed_full_cudagraphs_and_piecewise_cudagraphs

some names to consider instead of separate_attention_routine to communicate this flag is cudagraph related:
mixed_full_cudagraphs_and_piecewise_cudagraphs little long but leaves the door open to mixing full-cudagraphs with piecewise-cudagraphs if we want to dispatch on something other than is_pure_decode (see comment above)
full_cudagraphs_for_pure_decode_only little shorter and aligns better with current behavior but leaves little room to expand on this behavior in the future

I think we should get rid of is_pure_decode here. I could see cases in the future where we might want to mixed piecewise + full-cudagraphs but have a different heuristic than is_pure_decode. e.g. for FA3 we may want to run mixed small decodes or spec-decode using full-cudagraphs but run large prefills using piecewise cudagraphs

Honestly I find all these flags very confusing; I think I much simpler more extensible dispatch logic would be:

dispatch_key = DispatchKey(num_reqs=..., num_tokens=..., uniform_batch=...)
if dispatch_key in self.full_cudagraphs:
     return self.full_cudagraphs[dispatch_key]
# Fall-back if a full_cudagraph isn't available 
return self.piecewise_cudagraph or self.model

the uniform_batch flag here would indicate that all of the requests in the batch have the same number of tokens; so if num_reqs == num_tokens and uniform_batch would be pure decode but this would leave a door open for spec-decode support the future i.e. where "decode" steps are validating 2-4ish tokens at the same time. So if we are speculator is set to speculate 3 tokens at a time we could create full-cudagraphs for 3*num_reqs == num_tokens and uniform_batch. Something like FlashMLA would actually support this since the main thing it wants is a uniform batch.

Post from previous discussions.

You are right about more explicit control of full cudagraph like all, mixed, and decode_only. But I think an extra flag may not be necessary now, after the cudagraph_mode proposed in #20283, the extra control can be achieved by extending the enum mode.

Like use NONE=0, PIECEWISE=1, FULL=2, FULL_DECODE_ONLY=3, and FULL_AND_PIECEWISE=4. Here, NONE is for no cudagraph. PIECEWISE use only piecewise cudagraph (now v1 default). FULL means the current strategy for maximum full cudagraph support (with separate_attention_routine tunable to achieve mixed only or all). FULL_DECODE_ONLY uses only one set of cudagraph for pure decode, and no cudagraph for the rest. FULL_AND_PIECEWISE means explicitly having two sets of cudagraph, with full cudagraph for decode-only, and piecewise cudagraph for mixed batches or any rest. In this way, the separate_attention_routine is forced to true in FULL_DECODE_ONLY and FULL_AND_PIECEWISE, and the cascade attention can also be supported in these two modes.

I think @ProExpertProg and I agree to leave this explicit control of mixed_full_cudagraphs_and_piecewise_cudagraphs to a follow-up PR to include a new cudagraph_mode like FULL_AND_PIECEWISE. Currently, the FULL mode just maximizes the support of full cudagraph with proper fallbacks, and separate_attention_routine takes effect only in this mode to tell if we want to retain a unified routine or use separate routines for different situations.

But have to admit that separate_attention_routine seems a bit redundant now, as it would be overridden when attention backends cudagraph support is PURE_DECODE_ONLY or ALWAYS_UNIFIED, and would not be overridden only when ALWAYS_SEPARATE.

# Fall-back if a full_cudagraph isn't available 
return self.piecewise_cudagraph or self.model

I don't think we can now fully separate the piecewise_cudagraph and the raw model easily, since they are integrated together by the torch.compile with vllm piecewise compilation currently.

the uniform_batch flag here would indicate that all of the requests in the batch have the same number of tokens; so if num_reqs == num_tokens and uniform_batch would be pure decode but this would leave a door open for spec-decode support the future i.e. where "decode" steps are validating 2-4ish tokens at the same time. So if we are speculator is set to speculate 3 tokens at a time we could create full-cudagraphs for 3*num_reqs == num_tokens and uniform_batch. Something like FlashMLA would actually support this since the main thing it wants is a uniform batch.

I do agree we should leave a door for spec-decode, but I also think using the uniform_batch flag and num_tokens together for dispatching is somewhat confusing. First things is, if the speculator is set to speculate 3 tokens at a time, I guess this pattern is fixed, and we should just design a new enum representing that we are doing the speculate decode, rather than judging if 3*num_reqs == num_tokens and uniform_batch. Also, considering the meaning of uniform_batch you mentioned, that is not equivalent to pure decode and spec-decode. One counterexample may be in a pure prefill case where each request has 3 tokens.

Moreover, could we leave num_tokens being handled by the cudagraph wrapper itself? I think dispatching the num_tokens explicitly inside the dispatcher may be reasonable, but it is not trivial to manage all fused cases, and we couldn't reuse the uniform-designed cudagraph wrapper here. Leaving all that exact cudagraph management of one edge case to the wrapper would be good, as one cudagraph wrapper could just represent one case we want to process.

+1 to leaving the token dispatching inside the wrapper.

I don't think we can now fully separate the piecewise_cudagraph and the raw model easily, since they are integrated together by the torch.compile with vllm piecewise compilation currently.

this was just there as pseudo code to represent the logic; like you said the actual code would just be the model since the piecewise dispatching is under that (and dispatches to eager if the batch is too large)

One counterexample may be in a pure prefill case where each request has 3 tokens.

that would be ok; there would be no difference between that and a spec decode from the attention perspective so we would want it to use the the full-cudagraph if available in this case. We shouldn't fixate on the prefill/decode naming since chunked prefill and spec-decode blur these lines quite a bit; those names are just colloquially used to mean large query length or near 1 query length respectively. They are useful in the conversation but im a bit hesitant to harden those definitions (i.e. decode being query_len == 1) into the code; especially inside this fairly core part of the code.

Moreover, could we leave num_tokens being handled by the cudagraph wrapper itself?
+1 to leaving the token dispatching inside the wrapper.

I like the idea of having it as part of the dispatch keys because it leaves the door for have full-cudagraphs for only small batch sizes; like I could see a world where might only compile full cuda-graphs for up to BS 128 and then use piecewise for everything larger

But we already kinda support capturing full-cudagraphs only up to a size smaller than max_num_seqs via the configurable compilation_config.cudagraph_capture_sizes. So if that is the case then here would dispatch to the full-cudagraph wrapper only for it to fallback on eager inside the wrapper. I think this very confusing for readers/users. I think we should try to use this opportunity to flatten and simplify the dispatching; as someone who is moderately familiar with cudagraphs I vLLM I find all this dispatching very confusing.

This is where having a registry of the cudagraphs captured and keys representing the workloads they support I think could be far less confusing i.e. we can try to dispatch to a full-cudagraph first if one doesnt exist dispatch to a piecewise-cudagraph and if one doesnt exist dispatch to eager (I am aware the last 2 steps are currently happening in piecewise backend)

While I do see how later we might want to do more complex dispatching and size will certainly have to be involved, I think that's out of scope for this PR.

ya I don't think we have to go that far in this PR I just really want to make sure we are creating the right extensible abstractions since this PR is introducing alot of code that would have to get refactored to enable this

I just really want to make sure we are creating the right extensible abstractions since this PR

I agree with that & thanks for driving this point. What about the following "compromise"/alternate solution: the runtime style (or mode if we consolidate) and dispatch key dispatch separately. The runtime style is used to select between the CUDAGraphWrapper instances (full and piecewise), and the DispatchKey dispatches between recorded cudagraphs within a wrapper, for now only including num_tokens and uniform_batch. This way

Structure would look like this:

• CUDAGraphDispatcher decides the "runtime style" and the DispatchKey (whatever new name it receives), and sets it in the forward context
• Each CUDAGraphWrapper inspects the forward context and only captures/replays if runtime style matches. It uses the DispatchKey to decide which cudagraph to capture/replay.

This solves the issue where we can't do cudagraph capture/replay directly in the dispatcher for piecewise cudagraphs. While piecewise cudagraphs might not need as much detailed dispatching, this would give us the flexibility to dispatch to different routines for ops that were not in splitting ops.

@fhl2000 no need to implement this yet, let's reach a consensus on this first

With this approach, the CUDAGraphWrapper instances could have even less logic and blindly trust the forward context on what cg to dispatch to; if it doesn't exist yet, it gets captured. That way the dispatcher is the single source of truth on available cudagraphs.

If I'm missing something, a "single source of truth" for "available" cudagraphs a new noinit dictionary on CompilationConfig, used by both the dispatcher and the CUDAGraphWrapper instances.

Structure would look like this:

• CUDAGraphDispatcher decides the "runtime style" and the DispatchKey (whatever new name it receives), and sets it in the forward context
• Each CUDAGraphWrapper inspects the forward context and only captures/replays if runtime style matches. It uses the DispatchKey to decide which cudagraph to capture/replay.

Looks good to me.

With this approach, the CUDAGraphWrapper instances could have even less logic and blindly trust the forward context on what cg to dispatch to; if it doesn't exist yet, it gets captured. That way the dispatcher is the single source of truth on available cudagraphs.

If I'm missing something, a "single source of truth" for "available" cudagraphs a new noinit dictionary on CompilationConfig, used by both the dispatcher and the CUDAGraphWrapper instances.

Sharing a new noinit dictionary on CompilationConfig by both the dispatcher and the CUDAGraphWrapper instances seems unviable to me (or should be improved). While sharing DispatchKey on FULL style is possible as one graph item is enough for one dispatch key, it is not good for piecewise cudagraph as there are many graph items (almost one for a layer) corresponding to one DispatchKey in this case.

Yeah sorry phrased poorly/dropped a sentence. I was saying we could use that dictionary if we need a single source of truth but hopefully we don't. I don't think dispatch key cares about any layer wise info so the different piecewise backends (one per subgraph) don't need to worry about it.

Why do we have both?

I am pretty confused about the CUDAGraphMode vs CUDAGraphRuntimeStyle. Is the reason that CUDAGraphMode=Full for FAv2 means that we do FULL for prefill but PIECEWISE for decode?

If so, would anyone want to do piecewise for prefill and piecewise for decode for FAv2?

+1 why do we have 2 identical enums

Why do we have both?
+1 why do we have 2 identical enums

Sorry for causing confusion. But as explained in the comments in the code and also in the history chats, the CUDAGraphMode is intended for configure the behavior of modes. This mode configuration would be extended in future PR to contain more, for example FULL_ONLY_DECODE, FULL_AND_PIECEWISE, or something like AUTO as mentioned in #20283 that automatically selects among the previous modes. See below for the meaning of these modes

Like use NONE=0, PIECEWISE=1, FULL=2, FULL_DECODE_ONLY=3, and FULL_AND_PIECEWISE=4. Here, NONE is for no cudagraph. PIECEWISE use only piecewise cudagraph (now v1 default). FULL means the current strategy for maximum full cudagraph support (with separate_attention_routine tunable to achieve mixed only or all). FULL_DECODE_ONLY uses only one set of cudagraph for pure decode, and no cudagraph for the rest. FULL_AND_PIECEWISE means explicitly having two sets of cudagraph, with full cudagraph for decode-only, and piecewise cudagraph for mixed batches or any rest.

On the other side, CUDAGraphRuntimeStyle would be the actual style of cudagraphs we selected to run at runtime. I think there are only three styles to be shared among all possibilities, and can't be extended. This is also used as a property assigned for the CUDAGraph wrapper class for correctly activating cudagraphs of the right style, because currently we could have nested CUDAGraph wrappers, i.e., piecewise cudagraph wrapper integrated with the piecewise compiled model inside, while one wrapper wrapped outside for full cudagraph.

They have members of the same name coincidentally now, but should look fine after extending the cudagraph mode.

I still think mode can be enough and we just add asserts that during runtime we're not using any of the "mixed" modes that aren't valid as runtime styles.

I am pretty confused about the CUDAGraphMode vs CUDAGraphRuntimeStyle. Is the reason that CUDAGraphMode=Full for FAv2 means that we do FULL for prefill but PIECEWISE for decode?

No. For FA2, when CUDAGraphMode is set to Full, it means using FULL cudagraphs for both mixed prefill-decode stages and pure decode stages. However, since FA2's cudagraph support is marked as ALWAYS_SEPARATE, it prefers separate cudagraph routines for these two stages. Only when the separate_attention_routine is set to False, there will be a single FULL cudagraph for mixed prefill-decode batches, which is also compatible with pure decode scenarios.

Honestly, the current design of FULL mode just intends to do the max it can to support full cudagraph runtime style, while falling back to piecewise cudagraph or no cudagraph runtime style if any incompatible routine.

I still think mode can be enough and we just add asserts that during runtime we're not using any of the "mixed" modes that aren't valid as runtime styles.

But I think fusing the usage of cudagraph mode for both the semantics of "mode" and "runtimestyle" would lead to more confusion.

But I think fusing the usage of cudagraph mode for both the semantics of "mode" and "runtimestyle" would lead to more confusion.

I think it's a tradeoff. And I don't think the semantics are different enough. It's also easier to introduce a separate enum later if needed. To distinguish between them, all variables currently of type RuntimeStyle can have runtime_mode in the name instead of just name if you want the semantic meaning to be clearer.

@LucasWilkinson @zou3519 what do you think?